import os
import asyncio

from ragas.testset import TestsetGenerator
from ragas.testset.persona import Persona
from ragas.testset.synthesizers.single_hop.specific import (
    SingleHopSpecificQuerySynthesizer,
)
from common.utils import save_result, get_loader, get_maas_models, get_filename_without_ext, get_files_to_process
from common.config import config
from ragas.testset.synthesizers import QueryDistribution


async def adapte_language(language, generator_llm):
    distribution = [
        (SingleHopSpecificQuerySynthesizer(llm=generator_llm), 1.0),
    ]
    for query, _ in distribution:
        prompts = await query.adapt_prompts(language=language, llm=generator_llm)
        query.set_prompts(**prompts)
    return distribution


async def generate_testset(test_sets_file_path: str, docs: list) -> None:
    # 配置模型
    generator_llm, generator_embeddings = get_maas_models(
        base_url=config.maas.base_url,
        embedding_model=config.maas.embedding_model,
        model=config.maas.generator_model
    )
    # 设置角色
    personas = [
        Persona(
            name="好奇的学生",
            role_description="对世界充满好奇并希望更多地了解不同文化和语言的学生",
        ),
    ]
    # 创建测试集生成器
    generator = TestsetGenerator(
        llm=generator_llm, 
        embedding_model=generator_embeddings, 
        persona_list=personas)
    # 适应中文
    distribution: QueryDistribution = await adapte_language(language="chinese", generator_llm=generator_llm)
    # 生成测试集
    dataset = generator.generate_with_langchain_docs(
        docs, 
        testset_size=config.testset.size,
        query_distribution=distribution
    )
    # 保存结果
    save_result(dataset, test_sets_file_path)


async def main(specific_files: list[str] | None = None, test_sizes: list[int] = [10]):
    """
    生成测试集
    
    Args:
        specific_files: 指定要处理的文件列表，如果为None则处理docs_path目录下的所有文件
        test_sizes: 测试集大小列表，默认值为[10]
    """
    # 获取需要处理的文件列表
    print(f"文档目录: {config.path.docs_path}")
    files_to_process = get_files_to_process(specific_files, config.path.docs_path)
    if not files_to_process:
        return
    print(f"将处理以下 {len(files_to_process)} 个文件:")
    for file in files_to_process:
        print(f"  - {os.path.basename(file)}")

    # 循环测试不同大小的测试集
    for test_size in test_sizes:
        print(f"\n开始生成大小为 {test_size} 的测试集")
        config.testset.size = test_size
        # 处理文件
        for file_path in files_to_process:
            # 拼接测试集存放地址
            test_sets_file_path = os.path.join(
                config.path.results_path, 
                f"{get_filename_without_ext(file_path)}_QA_{config.testset.size}.xlsx"
            )
            try:
                # 获取文件加载器
                loadered_doc = get_loader(file_path).load()
                print(f"{file_path}加载完成, 开始生成测试集")
                # 生成测试集
                await generate_testset(test_sets_file_path, loadered_doc)
            except ValueError as e:
                print(f"跳过文件 {file_path}: {str(e)}")


if __name__ == "__main__":
    # 设置测试集大小列表
    test_sizes = list(range(70, 81, 10))  # [10, 20, 30, 40, 50, 60, 70, 80]
    # 指定要处理的文件列表
    specific_files = None
    asyncio.run(main(specific_files, test_sizes))

